import numpy as np
import torch
from matplotlib import pyplot as plt
from io import BytesIO


def show_points(coords, labels, ax, marker_size=375):
    pos_points = coords[labels == 1]
    neg_points = coords[labels == 0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white',
               linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white',
               linewidth=1.25)


def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)


def get_points_image(image, points_pred):
    fig, ax = plt.subplots()
    _,height, width = image.shape
    image = image.numpy().transpose(1, 2, 0)
    ax.imshow(image, cmap='gray')  # Display the grayscale image

    coords = points_pred[:, :2]
    coords = coords * np.array([width, height])
    labels = points_pred[:, 2]
    input_label = (labels > 0.5).int().reshape(-1, )
    show_points(coords, input_label, ax)

    # Additional steps to remove axis, ticks, and borders
    plt.axis('off')
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0, wspace=0, hspace=0)  # Remove any space around the plot
    fig.canvas.draw()  # Draw on canvas

    # Convert the drawn figure back to numpy array
    image_with_points = np.array(fig.canvas.renderer._renderer)
    plt.close(fig)  # Close the figure to free up memory

    return image_with_points[:, :, :3]  # Return the RGB image


def get_mask_image(image, mask, points_pred):
    fig, ax = plt.subplots()
    _,height, width = image.shape
    image = image.numpy().transpose(1, 2, 0)
    ax.imshow(image, cmap='gray')  # Display the grayscale image
    coords = points_pred[:, :2]
    coords = coords * np.array([width, height])
    show_mask(mask, plt.gca())
    # 将标签二进制化
    labels = points_pred[:, 2]
    input_label = (labels > 0.5).int().reshape(-1, )
    show_points(coords, input_label, ax)
    # Additional steps to remove axis, ticks, and borders
    plt.axis('off')
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0, wspace=0, hspace=0)  # Remove any space around the plot
    fig.canvas.draw()  # Draw on canvas

    # Convert the drawn figure back to numpy array
    image_with_points = np.array(fig.canvas.renderer._renderer)
    plt.close(fig)  # Close the figure to free up memory

    return image_with_points[:, :, :3]  # Return the RGB image
